package pers.whj.security;

import java.io.*;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.util.ArrayList;
import java.util.LinkedList;
import java.util.List;

/**
 * Use AES algorithm to encrypt files
 * encrypt one file will get two files(.wef and .wei)
 */
public class FileEncryption {
    private static final int LENGTH = 1024 * 1024;
    private static final String ENCRYPTION_SUFFIX = ".wef";
    private static final String INDEX_SUFFIX = ".wei";
    // avoid gen the same name for files to encrypt
    private static long sLastGenName = 0L;

    /**
     * interface to get the progress of the task
     */
    public interface Listener {
        /**
         * call when prepared
         *
         * @param length the total length, in bytes
         */
        void onPrepare(long length);

        /**
         * the current progress
         *
         * @param length the current length, in bytes
         */
        void onProcess(long length);

        /**
         * call when finished
         */
        void onFinish();
    }

    /**
     * encrypt a file or all the files in the directory
     *
     * @param path      path of file or directory
     * @param password  password string
     * @param recursive if path refer to a directory, encrypt all the files in its subdirectories
     *                  if it is true, or false only encrypt the files in this directory
     * @param delete    if delete the previous files
     * @param listener  the listener
     */
    public static void encrypt(String path, String password, boolean recursive, boolean delete
            , Listener listener) throws FileEncryptionException {
        File file = new File(path);
        if (!file.exists()) {
            throw new FileEncryptionException("file not exists!");
        }
        List<File> files = checkAndAddFiles(file, recursive, true);
        long total = 0L;
        if (listener != null) {
            for (File tmp : files) {
                total += tmp.length();
            }
            listener.onPrepare(total);
        }
        total = 0L;
        for (File tmp : files) {
            encryptFile(tmp, password, listener, total);
            total += tmp.length();
            if (delete && !tmp.delete()) {
                throw new FileEncryptionException("unsuccessfully delete the file!");
            }
        }
        if (listener != null) {
            listener.onFinish();
        }
    }

    /**
     * decrypt a file or all the files in the directory
     *
     * @param path      path of file or directory
     * @param password  password string
     * @param recursive if path refer to a directory, decrypt all the files in its subdirectories
     *                  if it is true, or false only decrypt the files in this directory
     * @param delete    if delete the previous files
     * @param listener  the listener
     */
    public static void decrypt(String path, String password, boolean recursive, boolean delete
            , Listener listener) throws FileEncryptionException {
        File file = new File(path);
        if (!file.exists()) {
            throw new FileEncryptionException("file not exists!");
        }
        List<File> files = checkAndAddFiles(file, recursive, false);
        long total = 0L;
        if (listener != null) {
            for (File tmp : files) {
                total += tmp.length();
            }
            listener.onPrepare(total);
        }
        total = 0L;
        for (File tmp : files) {
            decryptFile(tmp, password, listener, total);
            total += tmp.length();
            if (delete) {
                File indexFile = new File(tmp.getParent() + File.separator +
                        tmp.getName().substring(0, tmp.getName().length() - ENCRYPTION_SUFFIX.length()) + INDEX_SUFFIX);
                if (!tmp.delete() || !indexFile.delete()) {
                    throw new FileEncryptionException("unsuccessfully delete the file!");
                }
            }
        }
        if (listener != null) {
            listener.onFinish();
        }
    }

    /**
     * encrypt one file
     *
     * @param file     file
     * @param password password
     */
    private static void encryptFile(File file, String password, Listener listener, long current) throws FileEncryptionException {
        String preName = file.getName();
        String newName;
        synchronized (FileEncryption.class) {
            long genName = System.currentTimeMillis();
            if (genName <= sLastGenName) {
                genName = sLastGenName + 1;
            }
            sLastGenName = genName;
            newName = genName + "";
        }
        File newFile = new File(file.getParent() + File.separator + newName + ENCRYPTION_SUFFIX);
        File indexFile = new File(file.getParent() + File.separator + newName + INDEX_SUFFIX);
        try {
            if (newFile.createNewFile() && indexFile.createNewFile()) {
                FileInputStream input = new FileInputStream(file);
                FileOutputStream output = new FileOutputStream(newFile);
                DataOutputStream outputIndex = new DataOutputStream(Files.newOutputStream(indexFile.toPath()));
                // write the name of file
                byte[] write = AES.encrypt(preName, password).getBytes(StandardCharsets.UTF_8);
                output.write(write, 0, write.length);
                outputIndex.writeInt(write.length);
                byte[] read = new byte[LENGTH];
                int l = 0;
                while ((l = input.read(read)) > 0) {
                    write = AES.encrypt(read, 0, l, password);
                    output.write(write, 0, write.length);
                    outputIndex.writeInt(write.length);
                    current += l;
                    if (listener != null) {
                        listener.onProcess(current);
                    }
                }
                input.close();
                output.close();
                outputIndex.close();
            }
        } catch (Exception e) {
            throw new FileEncryptionException(e);
        }
    }

    private static void decryptFile(File file, String password, Listener listener, long current) throws FileEncryptionException {
        String newName = System.currentTimeMillis() + "";
        File newFile = new File(file.getParent() + File.separator + newName);
        File indexFile = new File(file.getParent() + File.separator +
                file.getName().substring(0, file.getName().length() - ENCRYPTION_SUFFIX.length()) + INDEX_SUFFIX);
        try {
            if (newFile.createNewFile()) {
                byte[] read = new byte[LENGTH * 2];
                int l = 0;
                String preName = null;
                try (FileInputStream input = new FileInputStream(file); DataInputStream inputIndex
                        = new DataInputStream(Files.newInputStream(indexFile.toPath()));
                     FileOutputStream output = new FileOutputStream(newFile)) {
                    int k = 1;
                    // read pre name
                    l = inputIndex.readInt();
                    k = input.read(read, 0, l);
                    byte[] write = AES.decrypt(read, 0, l, password);
                    preName = new String(write);
                    while (k > 0) {
                        l = inputIndex.readInt();
                        k = input.read(read, 0, l);
                        write = AES.decrypt(read, 0, l, password);
                        output.write(write, 0, write.length);
                        current += l;
                        if (listener != null) {
                            listener.onProcess(current);
                        }
                    }
                } catch (EOFException e) {
                    // rename the file
                    if (preName != null) {
                        newFile.renameTo(new File(file.getParent() + File.separator + preName));
                    }
                }
            }
        } catch (Exception e) {
            throw new FileEncryptionException(e);
        }
    }

    /**
     * check the target file and list all the files need to encrypt or decrypt
     *
     * @param file      target file
     * @param recursive if encrypt or decrypt the files in its subdirectories
     * @param encrypt   list files to encrypt if true, or to decrypt if false
     * @return file list
     */
    private static List<File> checkAndAddFiles(File file, boolean recursive, boolean encrypt) {
        List<File> files = new ArrayList<>();
        if (file.isFile() && isTargetFile(file, encrypt)) {
            files.add(file);
        }
        if (file.isDirectory()) {
            LinkedList<File> queue = new LinkedList<>();
            queue.add(file);
            do {
                File dir = queue.poll();
                File[] sub = dir.listFiles();
                if (sub != null) {
                    for (File tmp : sub) {
                        if (tmp.isFile() && isTargetFile(tmp, encrypt)) {
                            files.add(tmp);
                        }
                        if (tmp.isDirectory()) {
                            queue.offer(tmp);
                        }
                    }
                }
            } while (recursive && !queue.isEmpty());
        }
        return files;
    }

    private static boolean isTargetFile(File file, boolean encrypt) {
        return (encrypt && !isEncryptionFile(file))
                || (!encrypt && isEncryptionFile(file) && file.getName().endsWith(ENCRYPTION_SUFFIX));
    }

    private static boolean isEncryptionFile(File file) {
        String name = file.getName();
        return name.endsWith(ENCRYPTION_SUFFIX) || name.endsWith(INDEX_SUFFIX);
    }
}
